"""
@author Erik M Volz
@date June 28 2011
@version 0.1.0

Generates and solves ODEs for a compartmental model. 
"""
__version__='0.1.0'




from pylab import *
from scipy.integrate import odeint
from scipy.interpolate import interp1d,BarycentricInterpolator
import copy, cPickle, pdb, csv, numpy, string, time, os, sys 
import networkx
import coalescentSimulator, likelihood


def generate_random_model__regular_digraph(m, out_d_birth, out_d_migration):
	"""
	Generates directed graph representation of a compartmental model. 
	
	Parameters
	----------
	m : int
		number of states
	out_d_birth : int
		Number of birth terms for each state. 
	out_d_migration : int
		Number of migration terms for each state. 
	
	Value
	-----
	tuple of two networkx.DiGraphs
		each arc represents birth or migration
		
	"""
	
	notdone = True
	while notdone:
		bG = networkx.DiGraph()
		mG = networkx.DiGraph()
		for node in range(m):
			for k in range(out_d_birth):
				alter = randint(0, m)
				while bG.has_edge(node, alter): # exclude double edges
					alter = randint(0, m)
				#
				bG.add_edge(node, alter)
			#
			for k in range(out_d_migration):
				alter = node
				while alter==node:
					alter = randint(0, m)
				#
				mG.add_edge(node, alter)
			#
		#
		# test for connectedness
		uG = networkx.DiGraph() #networkx.union(mG, bG)
		uG.add_edges_from(bG.edges())
		uG.add_edges_from(mG.edges())
		if networkx.is_strongly_connected(uG):
			notdone = False
		#
	#
	return bG, mG 
#

def generate_random_model__poisson(m, z, p):
	"""
	Generates directed graph representation of a compartmental model. 
	
	Parameters
	----------
	m : int
		number of states
	z : float
		Expected number of birth or migration terms for each state. 
	p : float
		Probability that a link between states is birth and not migration. 
	
	Value
	-----
	tuple of two networkx.DiGraphs
		each arc represents birth or migration
		
	"""
	notdone = True
	while notdone:
		bG = networkx.DiGraph()
		mG = networkx.DiGraph()
		for node in range(m):
			bG.add_node(node)
			mG.add_node(node)
			for k in range(poisson(p * z)):
				alter = randint(0, m)
				while bG.has_edge(node, alter): # exclude double edges
					alter = randint(0, m)
				#
				bG.add_edge(node, alter)
			#
			for k in range(poisson((1-p)* z)):
				alter = node
				while alter==node:
					alter = randint(0, m)
				#
				mG.add_edge(node, alter)
			#
		#
		# test for connectedness
		uG = networkx.DiGraph() #networkx.union(mG, bG)
		uG.add_edges_from(bG.edges())
		uG.add_edges_from(mG.edges())
		if networkx.is_strongly_connected(uG):
			notdone = False
		#
	#
	return bG, mG #, networkx.adj_matrix(bG), networkx.adj_matrix(mG)
#

def generate_random_model__density(m, n_births, n_birth_loops, n_migs):
	"""
	Generates directed graph representation of a compartmental model. 
	The model has the given numbers of birth and migration terms distributed randomly between states. 
	
	Parameters
	----------
	m : int
		number of states
	n_births : int
		Number of birth terms for entire system. 
	n_birth_loops : int
		Number of birth terms from a state to itself. 
	n_migs : int
		Number of migration terms for entire system. 
	
	Value
	-----
	tuple of two networkx.DiGraphs
		each arc represents birth or migration
		
	"""
	notdone = True
	while notdone:
		bG = networkx.DiGraph()
		mG = networkx.DiGraph()
		bG.add_nodes_from(range(m))
		mG.add_nodes_from(range(m))
		for k in range(n_births):
			ego = randint(0, m)
			alter = ego
			while alter==ego or bG.has_edge(ego, alter):
				alter = randint(0, m)
				ego = randint(0, m)
			#
			bG.add_edge(ego, alter)
		#
		for k in range(n_birth_loops):
			ego = randint(0, m)
			while bG.has_edge(ego, ego):
				ego = randint(0, m)
			#
			bG.add_edge(ego, ego)
		#
		for k in range(n_migs):
			ego = randint(0, m)
			alter = ego
			while alter==ego or mG.has_edge(ego,alter):
				alter = randint(0, m)
				ego = randint(0, m)
			#
			mG.add_edge(ego, alter)
		#
		# test for connectedness
		uG = networkx.DiGraph() #networkx.union(mG, bG)
		uG.add_edges_from(bG.edges())
		uG.add_edges_from(mG.edges())
		if networkx.is_strongly_connected(uG):
			notdone = False
		#
	#
	return bG, mG 
#

class CompartmentalModel:
	"""
	CompartmentalModel(bG, mG, beta_lb, beta_ub, N, time)
	
	Solves an ODE compartmental model with arbitrary number of states. 
	Birth terms are logistic; death terms are mass action. 
	The initial conditions are one individual in each state. 
	
	
	Parameters
	----------
	bG : networkx.DiGraph
		Directed graph representation of birth terms between model states. 
	mG : networkx.DiGraph
		Directed graph representation of migration terms between model states. 
	beta_lb, beta_ub : float
		Lower and upper bound for birth rates. 
	N : int
		Population size. 
	time : array-like, float
		Time axis. Population size will calculated for every point on this axis. 
	
	Value
	-----
	Y : list of array (dimension m)
		Population size over time for each state. 
	cumulative_births, cumulative_migrations : list (length T) of array-like (dimension mXm)
	Y_i : list (length m) of scipy.interpolate.interp1d
		The interpolated population size for each state. 
	
	Methods
	-------
	births_migrations_prevalence(max_cumulative_births=inf)
		Generates a sequence of birth and migration matrices that can be used with methods in likelihood.py. 
	plot_model_structure(self):
		Makes a nice graphviz style graph for the random model with birth 
		and migration DiGraphs bG, mG.
	"""
	
	def _make_birth_matrix():
		self.bA = bA = array( networkx.adj_matrix(bG) )
		self.mA = mA = array( networkx.adj_matrix(mG) ) 
		self.birth_matrix = bA * uniform(low = beta_lb, high = beta_ub, size = m**2.).reshape((m,m))
		#~ pdb.set_trace()
		self.migration_matrix = mA
	#
	
	def __init__(self, bG, mG, beta_lb, beta_ub, N, time, **parms): 
		self.bG = bG; self.mG = mG; self.beta_lb = beta_lb; self.beta_ub = beta_ub
		
		self._make_birth_matrix()
		
		m = self.m = len(self.birth_matrix)
		self.N = N
		self.t = self.time = time
		
		self.__dict__.update(**parms)
		self._initialize()
	#
	def _initialize(self):
		m = self.m
		time = self.time
		
		
		self.x0 = concatenate( ( ones(m), zeros(self.m**2.), zeros(self.m**2.) ))
		self.x = x = odeint(self.xdot, self.x0, time)
		self.Y = self.x.T[:self.m]
		self.cumulative_births = [xx.reshape((self.m, self.m)) for xx in x.T[self.m:(self.m+self.m**2.)].T ]
		self.cumulativeTotalBirths = array( [ sum(xx) for xx in self.cumulative_births])
		self.cumulative_migrations = [xx.reshape((self.m, self.m)) for xx in x.T[(self.m+self.m**2.):].T ]
		
		self.x_i = [interp1d(self.time, xx) for xx in x.T]
		self.Y_i = self.x_i[:self.m] #[interp1d(self.time, YY) for YY in self.Y]
		try:
			self.cumulative_births_i = array( self.x_i[self.m:(self.m+self.m**2)]  ).reshape((self.m, self.m))
		except:
			pdb.set_trace()
		self.cumulative_migrations_i = array(self.x_i[(self.m+self.m**2):] ).reshape((self.m, self.m)) 
	#
	
	def xdot(self, x, t):
		Y = x[:self.m]
		sY = sum(Y)
		cumulative_births = x[self.m:(self.m+self.m**2.)].reshape((self.m, self.m))
		cumulative_migration = x[(self.m+self.m**2.):].reshape((self.m, self.m))
		d_cumulative_births = zeros((self.m, self.m))
		d_cumulative_migration = zeros((self.m, self.m))
		dx = zeros(self.m)
		for i in range(self.m):
			try:
				dx[i] += sum( ravel(self.birth_matrix.T[i]) * Y * (1. - sY/self.N) ) 
			except:
				pdb.set_trace()
			try:
				dx[i] += sum( ravel(self.migration_matrix.T[i]) * Y )
			except:
				pdb.set_trace()
			dx[i] -= sum( ravel(self.migration_matrix[i]) * Y[i] )
			dx[i] -= Y[i] #natural mortality
			
			for j in range(self.m):
				d_cumulative_births[j,i]    +=  Y[j] * self.birth_matrix[j,i] * (1. - sY/self.N) 
				d_cumulative_migration[j,i] +=  Y[j] * self.migration_matrix[j,i] 
		#
		
		return concatenate( (dx, ravel(d_cumulative_births), ravel(d_cumulative_migration) ) )
	#
	
	#axis_length = 1000 #length will be the same as the self.time variable
	def births_migrations_prevalence(self, max_cumulative_births=inf):
		"""
		Generates a sequence of birth and migration matrices that can be used with methods in likelihood.py. 
		
		Parameters
		----------
		max_cumulative_births : float
			If given, values will only be generated up to the point in 
			the simulation where the cumulative number of births does 
			not exceed this parameter. This can be useful for calibratin
			the output of this deterministic model and stochastic simulations. 
		
		Generates class variables
		-------------------------
		self.births, self.migrations : list of array-like (dimension mXm)
		self.prevalence : list of array (dimension m)
		"""
		if not isinf(max_cumulative_births):
			time_maxCumTransm = self.time[argmin(abs( max_cumulative_births-self.cumulativeTotalBirths ))]	
		else:
			time_maxCumTransm = self.time[-1]
		#
		self.axis_length = len(self.time)
		self.taxis =  linspace(0., time_maxCumTransm, self.axis_length+1)[1:]
		taxis = linspace(0., time_maxCumTransm, self.axis_length+1)
		
		self.births = list()
		self.migrations = list()
		self.prevalence = list()
		for t0, t1 in c_[taxis[:-1], taxis[1:]]:
			self.births.append( array( [ self.cumulative_births_i[k,l](t1) - self.cumulative_births_i[k,l](t0) for k in range(self.m) for l in range(self.m)]  ).reshape((self.m, self.m)) )
			
			self.migrations.append( array( [ self.cumulative_migrations_i[k,l](t1) - self.cumulative_migrations_i[k,l](t0) for k in range(self.m) for l in range(self.m)]  ).reshape((self.m, self.m)) )
			
			self.prevalence.append( array([ (self.Y_i[k](t1) + self.Y_i[k](t0))/2. for k in range(self.m) ] ) )
	#
	
	def plot_model_structure(self):
		"""
		makes a nice graphviz style graph for the random model with birth 
		and migration DiGraphs bG, mG;
		"""
		import  pygraphviz
		
		maxpenwidth = max(flatten(self.birth_matrix))
		minpenwidth = min(flatten(self.birth_matrix))
		deltapenwidth = maxpenwidth - minpenwidth
		
		G = pygraphviz.AGraph(strict = False, directed = True)
		G.graph_attr.update(dpi='600')
		for node in range(self.m):
			G.add_node(string.uppercase[node]) #, label = string.uppercase[node])
		for edge in self.bG.edges():
			G.add_edge(string.uppercase[edge[0]], string.uppercase[edge[1]])#, color='blue', penwidth=self.birth_matrix[edge[0], edge[1]])
			G.get_edge(string.uppercase[edge[0]], string.uppercase[edge[1]]).attr['color'] = 'blue'
			G.get_edge(string.uppercase[edge[0]], string.uppercase[edge[1]]).attr['penwidth'] = repr(1.+5*(self.birth_matrix[edge[0], edge[1]] - minpenwidth) / deltapenwidth)
		#
		#~ G.add_edges_from(bG.edges(), color = 'blue', penwidth=2.)
		#~ G.add_edges_from(mG.edges(), color = 'red', penwidth=1.)
		for edge in self.mG.edges():
			G.add_edge(string.uppercase[edge[0]], string.uppercase[edge[1]] )#
			G.get_edge(string.uppercase[edge[0]], string.uppercase[edge[1]]).attr['color'] = 'red'
		#
		
		G.node_attr.update(color='black', filled =repr( False))
		
		G.draw('figures/RandomModel__plot_model_structure.png', prog='circo')#args='-dpi=600')
		G.draw('figures/RandomModel__plot_model_structure.ps', prog='circo')
		G.draw('figures/RandomModel__plot_model_structure.svg', prog='circo')
	#
#

class CompartmentalModel2(CompartmentalModel):
	"""july 22 Extends CompartmentalModel, but birth and migration rates are not random variables. Rather they are given when the class is initialized. 
	"""
	def __init__(self,  bA, mA, N, time, **parms): 
		self.bA = bA; self.mA = mA; 
		self.birth_matrix=copy.deepcopy(bA)
		self.migration_matrix =copy.deepcopy(mA)
		m = self.m = len(self.birth_matrix)
		self.N = N
		self.t = self.time = time
		
		self.__dict__.update(**parms)
		self._initialize()
#


if __name__=='__main__':
	try:
		m = int(sys.argv[1])
	except:
		print 'usage: python modelGenerator.py numberOfStates'
	#
	N=10**4
	n_births = m; n_migs = m; n_birth_loops = 0
	bg, mg = generate_random_model__density(m, n_births, n_birth_loops, n_migs)
	ba = 2* array(networkx.adj_matrix(bg))
	ma = array(networkx.adj_matrix(mg))
	
	time = linspace(0., 15., 100)
	cm = CompartmentalModel2(ba, ma, N, time)
	
	
	
	cm.births_migrations_prevalence()
	wb = csv.writer(open('births_'+sys.argv[1] + '.csv', 'w'))
	wm = csv.writer(open('migrations_'+sys.argv[1] + '.csv', 'w'))
	wp = csv.writer(open('populationSize_'+sys.argv[1] + '.csv', 'w'))
	wt = csv.writer(open('time_'+sys.argv[1] + '.csv', 'w'))
	for b,m,p,t in zip(cm.births,cm.migrations, cm.prevalence, time):
		try:
			wb.writerows(b)
			wm.writerows(m)
			wp.writerows([p])
			wt.writerows([[t]])
		except:
			pdb.set_trace()
	#
	
	plot(cm.time, cm.Y.T)
	show()
